import cv2
import numpy as np
from ImageCorrector import ImageCorrector
from Camera import Camera
from DisparityGenerator import DisparityGenerator




class DisparityCutGenerator:
    def __init__(self):
        pass

    def ParallaxSegmentation(self):
        """
        将有堆料和无堆料的视差图像做差，视差小于14（不同照片的值不一样）的像素值置为0
        :return: 分割出堆料的视差图
        """
        file_path = "./相机信息/out.xls"
        imageCorrector28 = ImageCorrector("images/IMG_0028.jpg", file_path)
        imageCorrector27 = ImageCorrector("images/IMG_0027.jpg", file_path)
        # 左右视图转化为灰度，图像还是那个图像
        imgL = cv2.cvtColor(imageCorrector27.left_rectified, cv2.COLOR_BGR2GRAY)
        imgR = cv2.cvtColor(imageCorrector27.right_rectified, cv2.COLOR_BGR2GRAY)
        # 高斯模糊
        imgL = cv2.GaussianBlur(imgL, (3, 3), 0)
        imgR = cv2.GaussianBlur(imgR, (3, 3), 0)
        #生成视差
        disparity27 = self.compute_disparity(imgL, imgR, num=5, blockSize=15)

        imgL = cv2.cvtColor(imageCorrector28.left_rectified, cv2.COLOR_BGR2GRAY)
        imgR = cv2.cvtColor(imageCorrector28.right_rectified, cv2.COLOR_BGR2GRAY)

        imgL = cv2.GaussianBlur(imgL, (3, 3), 0)
        imgR = cv2.GaussianBlur(imgR, (3, 3), 0)

        disparity28 = self.compute_disparity(imgL, imgR, num=5, blockSize=15)
        """
        结果 diff 也是一个 图像（ndarray），表示 disparity27 和 disparity28 之间的差异程度

        mask 就是差异大于一定程度的蒙版
        """
        diff = cv2.absdiff(disparity27, disparity28)

        # 处理深度差异（小于 threshold_value 的置为 0）
        threshold_value = 14
        diff[diff < threshold_value] = 0
        mask = diff > 0  # 选出大于 0 的像素

        new_disparity28 = (DisparityGenerator(ImageCorrector("images/IMG_0028.jpg", file_path)).generate(5, 15, False))[
            0]#生成有堆料图的视差图

        # 抠图，将视差变化不大区域置为0
        disparity28_filtered = np.where(mask, new_disparity28, 0.0)
        return disparity28_filtered

    def compute_disparity(self,imgL, imgR, num=5, blockSize=7):
        """
        计算视差图并进行WLS滤波
        return：生成的视差图，经过处理，和DisparityGenerator类中的稍微有点不同，大致相同
        """
        if blockSize % 2 == 0:
            blockSize += 1  # 确保为奇数

        stereo_params = {
            'minDisparity': 0,
            'numDisparities': 16 * num,
            'blockSize': blockSize,
            'P1': 8 * 3 * blockSize ** 2,
            'P2': 32 * 3 * blockSize ** 2,
            'disp12MaxDiff': 1,
            'uniquenessRatio': 10,
            'speckleWindowSize': 100,
            'speckleRange': 32
        }

        stereo = cv2.StereoSGBM_create(**stereo_params)
        right_matcher = cv2.ximgproc.createRightMatcher(stereo)

        disparity_left = stereo.compute(imgL, imgR).astype(np.float32) / 16.0
        disparity_right = right_matcher.compute(imgR, imgL).astype(np.float32) / 16.0

        wls_filter = cv2.ximgproc.createDisparityWLSFilterGeneric(False)
        wls_filter.setSigmaColor(1.5)
        disparity_filtered = wls_filter.filter(disparity_left, imgL, None, disparity_right)
        return disparity_filtered
